#ifndef _PSO_H_
#define _PSO_H_

#include <vector>
#include <string>
#include <sstream>
#include <iostream>
#include <random>
#include <time.h>
#include <cstdlib>

#include "CImg.h"
using namespace cimg_library;

namespace Utils 
{

class Random
{
public:
	Random()
	{
		srand( (unsigned int)(time(NULL)) );
	}

	Random(unsigned int seed)
	{
		srand(seed);
	}

	int Next()
	{
		return rand();
	}

	double NextDouble()
	{
		return (double)(rand())/RAND_MAX;
	}

	int Next(int minValue, int maxValue)
	{
		return minValue + rand()%(maxValue - minValue);
	}

}; // End of class Random


}

namespace PSO {

class Particle
{
private:
	std::vector<double> position_;
	double              fitness_;
	std::vector<double> velocity_;

	std::vector<double> bestPosition_;
	double              bestFitness_;


public:
	Particle() { }
	Particle(
		const std::vector<double>& position,
		const double&              fitness,
		const std::vector<double>& velocity,
		const std::vector<double>& bestPosition,
		const double&              bestFitness
		) {
			for (unsigned int i=0; i<position.size(); i++)
				position_.push_back( position[i] );

			fitness_ = fitness;

			for (unsigned int i=0; i<velocity.size(); i++)
				velocity_.push_back( velocity[i] );

			for (unsigned int i=0; i<bestPosition.size(); i++)
				bestPosition_.push_back( bestPosition[i] );

			bestFitness_ = bestFitness;
	}
	~Particle(){};
	
	double get_fitness() const {
		return fitness_;
	}

	void copy_position_to(std::vector<double>& position) const {
		for (unsigned int i=0; i<position_.size(); i++)
			position[i] = position_[i];
	}

	void copy_position_from(const std::vector<double>& position) {
		for (unsigned int i=0; i<position_.size(); i++)
			position_[i] = position[i];
	}

	void copy_velocity_from(const std::vector<double>& velocity) {
		for (unsigned int i=0; i<velocity_.size(); i++)
			velocity_[i] = velocity[i];
	}

	void copy_best_position_from(const std::vector<double>& best_position) {
		for (unsigned int i=0; i<bestPosition_.size(); i++)
			bestPosition_[i] = best_position[i];
	}

	double get_velocity(unsigned int i) const {
		return velocity_[i];
	}

	double get_position(unsigned int i) const {
		return position_[i];
	}

	double get_best_position(unsigned int i) const {
		return bestPosition_[i];
	}

	double get_best_fitness() const {
		return bestFitness_;
	}

	void set_fitness(const double& fitness) {
		fitness_ = fitness;
	}

	void set_best_fitness(const double& best_fitness) {
		bestFitness_ = best_fitness;
	}

	void ToString() {
		std::stringstream ss;

		ss << "==============================\n";
		ss << "Position: ";
		for (unsigned int i=0; i<position_.size(); i++)
			ss << position_[i] << "   ";
		ss << "\n";
		ss << "Fitness : " << fitness_ << "\n";
		ss << "Velocity: ";
		for (unsigned int i=0; i<velocity_.size(); i++)
			ss << velocity_[i] << "   ";
		ss << "\n";
		ss << "Best Position: ";
		for (unsigned int i=0; i<bestPosition_.size(); i++)
			ss << bestPosition_[i] << "   ";
		ss << "\n";
		ss << "Best Fitness : " << bestFitness_ << "\n";
		ss << "===============================\n";

		std::cout << ss.str();
	}

}; // End of class Particle


double ObjectiveFunc(const std::vector<double>& X);



class Swarm 
{
private:
	std::vector<Particle*> swarm_;

	int Dim_;

	double w_, c1_, c2_, c3_;

	std::vector<double> BoundPosMin_;
	std::vector<double> BoundPosMax_;
	std::vector<double> BoundVeloMin_;
	std::vector<double> BoundVeloMax_;

	std::vector<double> bestGlobalPosition_;
	double              bestGlobalFitness_;

	std::vector<double> superGlobalPosition_;
	double              superGlobalFitness_;

	void init(Utils::Random& ran) {
		for (unsigned int p=0; p<swarm_.size(); p++) {
			std::vector<double> ranPos(Dim_);
			for (int i=0; i<Dim_; i++) {
				double lo = BoundPosMin_[i];
				double hi = BoundPosMax_[i];
				double r = ran.NextDouble();
				ranPos[i] = (hi - lo) * ran.NextDouble() + lo;
			}

			double fitness = ObjectiveFunc(ranPos);

			std::vector<double> ranVelo(Dim_);
			for (int i=0; i<Dim_; i++) {
				double lo = BoundVeloMin_[i];
				double hi = BoundVeloMax_[i];
				ranVelo[i] = (hi - lo) * ran.NextDouble() + lo;
			}

			swarm_[p] = new PSO::Particle(
				ranPos,fitness,ranVelo,ranPos,fitness);

			if (swarm_[p]->get_fitness() > bestGlobalFitness_) {
				bestGlobalFitness_ = swarm_[p]->get_fitness();
				swarm_[p]->copy_position_to(bestGlobalPosition_);
			}
		}
	}

public:
	Swarm() {}
	~Swarm() {
		for (unsigned int i=0; i<swarm_.size(); i++) {
			delete swarm_[i];
		}
	}

	Swarm(
		const int numParticles,
		const std::vector<double>& bndPosMin, 
		const std::vector<double>& bndPosMax, 
		const std::vector<double>& bndVeloMin, 
		const std::vector<double>& bndVeloMax,
		Utils::Random& ran) 
		:w_(0.729),c1_(1.49445),c2_(1.49445),c3_(1.49445)
	{

			Dim_ = (int)bndPosMin.size();
			BoundPosMin_.resize(Dim_);
			BoundPosMax_.resize(Dim_);
			BoundVeloMin_.resize(Dim_);
			BoundVeloMax_.resize(Dim_);

			for (int i=0; i<Dim_; i++) {
				BoundPosMin_[i] = bndPosMin[i];
				BoundPosMax_[i] = bndPosMax[i];
				BoundVeloMin_[i] = bndVeloMin[i];
				BoundVeloMax_[i] = bndVeloMax[i];
			}

			swarm_.resize(numParticles);
			bestGlobalPosition_.resize(Dim_);
			bestGlobalFitness_ = DBL_MIN;

			superGlobalPosition_.resize(Dim_);
			superGlobalFitness_ = DBL_MIN;


			init(ran);
	}

	void set_weights(double w, double c1, double c2, double c3) {
		w_ = w; c1_ = c1; c2_ = c2; c3_ = c3;
	}

	void Run(
		const int numIterations, Utils::Random& ran, 
		CImg<int>* region = NULL, 
		CImg<unsigned char>* seg_face = NULL,
		std::vector<std::pair<double,double>> * linechart_in = NULL,
		CImg<unsigned char>* img = NULL
		) 
	{

		CImgDisplay disp;
		CImg<unsigned char> canvas;
		const char red[] = {255,0,0}, blue[] = {0,0,255}, black[] = {0,0,0};

		if (img != NULL) {
			canvas.assign(*img);

			double tx,ty,sx,sy,r1,r2;
			for (unsigned int i=0; i<swarm_.size(); i++) {
				tx = swarm_[i]->get_position(0);
				ty = swarm_[i]->get_position(1);
				sx = swarm_[i]->get_position(2);
				sy = swarm_[i]->get_position(3);

				printf("%f, %f, %f, %f, %f, %f\n", tx, ty, sx, sy, tx*sx, ty*sy);
				tx *= sx;
				ty *= sy;

				if ((int)tx<0 || (int)tx>=img->width() || (int)ty<0 || (int)ty>=img->height())
					continue;

				r1 = (4.0*sx < 1.0) ? 1.0: 4.0*sx;
				r2 = (4.0*sy < 1.0) ? 1.0: 4.0*sy;
				canvas.draw_ellipse((int)tx, (int)ty, r1, r2, 0, red);
				
			}
		
			tx = bestGlobalPosition_[0];
			ty = bestGlobalPosition_[1];
			sx = bestGlobalPosition_[2];
			sy = bestGlobalPosition_[3];

			tx *= sx; 
			ty *= sy;

			if (!((int)tx<0 || (int)tx>=img->width() || (int)ty<0 || (int)ty>=img->height())) {
				r1 = (4.0*sx < 1.0) ? 1.0: 4.0*sx;
				r2 = (4.0*sy < 1.0) ? 1.0: 4.0*sy;			
				canvas.draw_ellipse((int)tx, (int)ty, r1, r2, 0, blue);
			}
			canvas.display(disp);
			disp.set_title("Running a swarm...");
		}

		std::vector<double> newVelo(Dim_);
		std::vector<double> newPos(Dim_);
		double newFitness, r1,r2,r3;
		
		for (int iter = 0; iter < numIterations; iter++) {
			for (unsigned int p=0; p<swarm_.size(); p++) {
				PSO::Particle* currP = swarm_[p];

				for (int j=0; j<Dim_; j++) {
					r1 = ran.NextDouble();
					r2 = ran.NextDouble();
					r3 = ran.NextDouble();

					double 
						local = currP->get_best_position(j) - currP->get_position(j),
						global = bestGlobalPosition_[j] - currP->get_position(j),
						super = superGlobalPosition_[j] - currP->get_position(j);

					newVelo[j] = 
						(w_*currP->get_velocity(j)) + 
						(c1_*r1*local) + (c2_*r2*global) + (c3_*r3*super);

					if (newVelo[j] < BoundVeloMin_[j])
						newVelo[j] = BoundVeloMin_[j];
					else if (newVelo[j] > BoundVeloMax_[j])
						newVelo[j] = BoundVeloMax_[j];
				}// Each j

				currP->copy_velocity_from(newVelo);

				for (int j=0; j<Dim_; j++) {
					newPos[j] = currP->get_position(j) + newVelo[j];
					if (newPos[j] < BoundPosMin_[j])
						newPos[j] = BoundPosMin_[j];
					else if (newPos[j] > BoundPosMax_[j])
						newPos[j] = BoundPosMax_[j];
				}

				currP->copy_position_from(newPos);

				newFitness = ObjectiveFunc(newPos);
				currP->set_fitness(newFitness);

				if (newFitness > currP->get_best_fitness()) {

					currP->copy_best_position_from(newPos);
					currP->set_best_fitness(newFitness);
				}

				if (newFitness > bestGlobalFitness_) {
					for (int i=0; i<Dim_; i++)
						bestGlobalPosition_[i] = newPos[i];
					bestGlobalFitness_ = newFitness;
				}
			} // Each Particle

			if (img != NULL) {
				canvas.assign(*img);
				cimg_forXYC(canvas,x,y,c) {
					if ((*region)(x,y)==1) continue;
					canvas(x,y,c) = canvas(x,y,c)*0.3f + (*seg_face)(x,y,c)*0.7f;
				}
	
				double tx,ty,sx,sy,r1,r2;
				for (unsigned int i=0; i<swarm_.size(); i++) {
					tx = swarm_[i]->get_position(0);
					ty = swarm_[i]->get_position(1);
					sx = swarm_[i]->get_position(2);
					sy = swarm_[i]->get_position(3);
	
					tx *= sx;
					ty *= sy;
	
					if ((int)tx<0 || (int)tx>=img->width() || (int)ty<0 || (int)ty>=img->height())
						continue;
	
					r1 = (4.0*sx < 1.0) ? 1.0: 4.0*sx;
					r2 = (4.0*sy < 1.0) ? 1.0: 4.0*sy;
					canvas.draw_ellipse((int)tx, (int)ty, r1, r2, 0, red);
				}
	
				tx = bestGlobalPosition_[0];
				ty = bestGlobalPosition_[1];
				sx = bestGlobalPosition_[2];
				sy = bestGlobalPosition_[3];
	
				tx *= sx;
				ty *= sy;
	
				if (!((int)tx<0 || (int)tx>=img->width() || (int)ty<0 || (int)ty>=img->height())) {
					r1 = (4.0*sx < 1.0) ? 1.0: 4.0*sx;
					r2 = (4.0*sy < 1.0) ? 1.0: 4.0*sy;			
					canvas.draw_ellipse((int)tx, (int)ty, r1, r2, 0, blue);
				}
	
				tx = bestGlobalPosition_[0];
				ty = bestGlobalPosition_[1];
				double x0 = ((*linechart_in)[0].first + tx)*sx;
				double y0 = ((*linechart_in)[0].second + ty)*sy;
				for (unsigned int i=1; i<(*linechart_in).size(); i++) {
					double x1 = ((*linechart_in)[i].first + tx)*sx;
					double y1 = ((*linechart_in)[i].second + ty)*sx;
	
					if ((int)x0<0 || (int)x0>=img->width() || (int)y0<0 || (int)y0>img->height() ||
						(int)x1<0 || (int)x1>=img->width() || (int)y1<0 || (int)y1>img->height()) {
							;
					} else {
						canvas.draw_line((int)x0,(int)y0,(int)x1,(int)y1,blue);
					}
					x0 = x1; y0 = y1;
				}
	
				std::stringstream ss;
				ss << iter << " / " << numIterations;
				canvas.draw_text(20,20,ss.str().c_str(),black,0);
	
				ss.str(std::string()); 
				ss << "Best Fit: " << bestGlobalFitness_;
				canvas.draw_text(20,40,ss.str().c_str(),black,0);
	
				ss.str(std::string());
				ss << "Best Sol: " << bestGlobalPosition_[0] << " , " << bestGlobalPosition_[1] << " , "
				                   << bestGlobalPosition_[2] << " , " << bestGlobalPosition_[3];
				canvas.draw_text(20,60,ss.str().c_str(),black,0);
	
				canvas.display(disp);
				disp.wait(100);
			}

		} // Main PSO loop

		if (img != NULL) {
			while (!disp.is_closed());
		}
	}

	double get_best_fitness() const {
		return bestGlobalFitness_;
	}

	void copy_best_position_to(std::vector<double>& pos) const {
		for (int i=0; i<Dim_; i++) {
			pos[i] = bestGlobalPosition_[i];
		}
	}

	void set_super_fitness(const double& superFitness) {
		superGlobalFitness_ = superFitness;
	}

	void copy_super_position_from(const std::vector<double>& pos) {
		for (int i=0; i<Dim_; i++) 
			superGlobalPosition_[i] = pos[i];
	}

	double get_best_position(int i) const {
		return bestGlobalPosition_[i];
	}
};



























} // End of PSO namespace



#endif

















